# Copyright (c) 2024-present AI-Labs

import time
import requests, json


from comps import (
    CustomLogger,
    TryonInputs,
    SDOutputs,
    ServiceType,
    opea_microservices,
    register_microservice,
    register_statistics,
    statistics_dict,
)

from configs import config

logger = CustomLogger("opea_service@tryon")
logflag = config.opea_service.tryon.logs

"""
注册微服务
"""
@register_microservice(
    name="opea_service@tryon",
    service_type=ServiceType.TRYON,
    host=config.opea_service.tryon.host,
    port=config.opea_service.tryon.port,
    endpoint="/v1/tryon",
    input_datatype=TryonInputs,
    output_datatype=SDOutputs,
)

# 微服务的具体处理逻辑
@register_statistics(names=["opea_service@tryon"])
def tryon(input: TryonInputs):
    start = time.time()

    person_image = input.person_image
    person_mask = input.person_mask
    cloth_image = input.cloth_image
    cloth_mask = input.cloth_mask
    num_inference_steps = input.num_inference_steps
    guidance_scale = input.guidance_scale
    seed = input.seed
    if logflag:
        logger.info(f"接收到用户请求")

    url = f"{config.opea_service.tryon.endpoint}/try_on/leffa/v1"

    headers = {'Content-Type': 'application/json; charset=utf-8'}
    data=json.dumps({
        'person_image': person_image,
        'person_mask': person_mask,
        'cloth_image': cloth_image,
        'cloth_mask': cloth_mask,
        'num_inference_steps': num_inference_steps,
        'guidance_scale': guidance_scale,
        'seed': seed
    })

    # 请求底层基础功能进行图片处理
    response = requests.post(url=url, headers=headers, data=data)

    # 统计耗时
    statistics_dict["opea_service@tryon"].append_latency(time.time() - start, None)

    if logflag:
        logger.info(f"接收到处理结果：{response}")

    # 返回响应结果
    return SDOutputs(images=response.json()["images"])

"""
启动微服务
"""
def start():
    opea_microservices["opea_service@tryon"].start()

if __name__ == "__main__":
    start()
